#!/usr/bin/perl

#####################
# DNS UDP/TCP stress
#####################

use strict;
use warnings;
use threads;
use threads::shared;
use IO::Socket;
use Term::ReadKey;

# default timeout
my $cl_timeout = 5;

if (@ARGV != 6 && @ARGV != 7) {
	print STDERR "Usage: $0 <p|s> <u|t> <target_ip> <target_port> <local_port|num_threads> <name> [<timeout=$cl_timeout>]\n";
	exit 1;
}

if (not $ARGV[0] eq 'p' and not $ARGV[0] eq 's') {
	print STDERR "mode must be either 'p' or 's'\n";
	exit 1;
}

if (not $ARGV[1] eq 'u' and not $ARGV[1] eq 't') {
	print STDERR "transport must be either 'u' or 't'\n";
	exit 1;
}

my $count = 0;
my @conns : shared;
my $needstop : shared;

if (@ARGV == 7) {
	$cl_timeout = $ARGV[6];
	if ($cl_timeout < 0) {
		$cl_timeout = 0;
	}
}

srand;

# transaction id
my $dnsreq = pack 'C2', int rand 256, int rand 256;

# header misc
$dnsreq .= pack 'C10', 0x01, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00;

# query name
my $reqname = $ARGV[5];

if (not substr($reqname, -1, 1) eq '.') {
     $reqname .= '.';
}

foreach (split(/\./, $reqname)) {
     $dnsreq .= pack 'C1', length ($_);
     $dnsreq .= $_;
}

# terminating zero
$dnsreq .= pack 'C1', 0;

# type A class IN
$dnsreq .= pack 'C4', 0x00, 0x01, 0x00, 0x01;

# heading length if TCP
if ($ARGV[1] eq 't') {
	$dnsreq = (pack 'n1', length ($dnsreq)) . $dnsreq;
}

print "** spawning DNS requests with max timeout=$cl_timeout, any key to cancel...\n";

$needstop = 0;

push @conns, &share(async {
	while ($needstop == 0) {
		my $client = IO::Socket::INET->new(
			Proto => $ARGV[1] eq 'u' ? 'udp' : 'tcp',
			LocalPort => $ARGV[0] eq 'p' ? $ARGV[4] : 0,
			PeerAddr => $ARGV[2],
			PeerPort => $ARGV[3],
			);
		if ($client) {
			$client->send($dnsreq);
			sleep 1 + int rand $cl_timeout;
			$client->shutdown(2);
			$client->close();
		}
	}
}) until $count++ == ($ARGV[0] eq 's' ? $ARGV[4] : 1);

ReadMode('cbreak');

while (!ReadKey(-1)) {
	threads->yield();
}

$needstop = 1;

ReadMode('normal');

print "** exiting...\n";

foreach (threads->list()) {
	$_->join();
}
